#pragma once

#include <mcp/config.h>
#include <mcp/forward.h>
#include <mcp/types.h>
#include <mcp/error.h>

#include <functional>
#include <memory>
#include <string>
#include <vector>
#include <map>
#include <optional>
#include <coroutine>

namespace mcp {
namespace server {

/// @brief Forward declarations
class ToolManager;
class ResourceManager;
class PromptManager;
class ServerSession;

namespace detail {
    template<typename ServerSessionT, typename LifespanT>
    class ContextImpl;
}

/// @brief Server configuration
struct ServerConfig {
    std::string name;
    std::string version = "1.0.0";
    std::string description;
    bool debug = false;
    std::string host = "localhost";
    uint16_t port = 8080;
    std::string mount_path = "/";
    std::string sse_path = "/sse";
    std::string streamable_http_path = "/mcp";
    bool warn_on_duplicate_resources = true;
    bool warn_on_duplicate_tools = true;
    bool warn_on_duplicate_prompts = true;
};

/// @brief Tool registration options
struct ToolOptions {
    std::optional<std::string> description;
    std::optional<json::Schema> input_schema;
    std::optional<json::Schema> output_schema;
    types::ToolAnnotations annotations;
};

/// @brief Resource registration options
struct ResourceOptions {
    std::optional<std::string> name;
    std::optional<std::string> description;
    std::string mime_type = "text/plain";
};

/// @brief Prompt registration options
struct PromptOptions {
    std::optional<std::string> description;
};

/// @brief Lifespan context for server lifecycle management
template<typename LifespanT>
class LifespanManager {
public:
    template<typename Func>
    void on_startup(Func&& func);

    template<typename Func>
    void on_shutdown(Func&& func);

    LifespanT& get_context() { return context_; }
    const LifespanT& get_context() const { return context_; }

private:
    std::vector<std::function<void()>> startup_handlers_;
    std::vector<std::function<void()>> shutdown_handlers_;
    LifespanT context_;
};

/// @brief FastMCPServer - High-level MCP server implementation
class FastMCPServer : public std::enable_shared_from_this<FastMCPServer> {
public:
    explicit FastMCPServer(const ServerConfig& config = {});
    ~FastMCPServer();

    // Server lifecycle
    template<typename TransportT, typename... Args>
    void run(Args&&... args);

    template<typename LifespanT, typename Func>
    void lifespan(Func&& func);

    // Tool registration
    template<typename Func>
    void tool(const std::string& name, Func&& func, const ToolOptions& options = {});

    // Resource registration
    template<typename Func>
    void resource(const std::string& uri_template, Func&& func, const ResourceOptions& options = {});

    // Prompt registration
    template<typename Func>
    void prompt(const std::string& name, Func&& func, const PromptOptions& options = {});

    // Server information
    const ServerConfig& config() const { return config_; }
    std::vector<types::Tool> list_tools() const;
    std::vector<types::Resource> list_resources() const;
    std::vector<types::ResourceTemplate> list_resource_templates() const;
    std::vector<types::Prompt> list_prompts() const;

    // Session management
    std::shared_ptr<ServerSession> create_session();

private:
    struct Impl;
    std::unique_ptr<Impl> impl_;
    ServerConfig config_;
    std::unique_ptr<ToolManager> tool_manager_;
    std::unique_ptr<ResourceManager> resource_manager_;
    std::unique_ptr<PromptManager> prompt_manager_;
    std::unique_ptr<LifespanManager<void>> lifespan_manager_;

    friend class detail::ContextImpl<ServerSession, void>;
};

/// @brief Context for tool/resource/prompt handlers
template<typename ServerSessionT = ServerSession, typename LifespanT = void>
class Context {
public:
    // Logging
    void debug(const std::string& message);
    void info(const std::string& message);
    void warning(const std::string& message);
    void error(const std::string& message);

    // Progress reporting
    void report_progress(double progress,
                        std::optional<double> total = std::nullopt,
                        std::optional<std::string> message = std::nullopt);

    // Resource access
    Result<types::ResourceContent> read_resource(const std::string& uri);

    // User interaction
    template<typename SchemaT>
    Result<types::ElicitationResult> elicit(const std::string& message, const SchemaT& schema);

    // LLM sampling
    Result<types::SamplingResult> create_message(
        const std::vector<types::SamplingMessage>& messages,
        const types::SamplingOptions& options = {});

    // Context access
    ServerSessionT& session();
    LifespanT& lifespan_context();

private:
    friend class FastMCPServer;
    friend class detail::ContextImpl<ServerSessionT, LifespanT>;

    Context(std::weak_ptr<ServerSessionT> session,
            std::shared_ptr<LifespanT> lifespan_context,
            types::RequestId request_id);

    std::weak_ptr<ServerSessionT> session_;
    std::shared_ptr<LifespanT> lifespan_context_;
    types::RequestId request_id_;
};

// Template implementations

template<typename TransportT, typename... Args>
void FastMCPServer::run(Args&&... args) {
    auto transport = std::make_unique<TransportT>(std::forward<Args>(args)...);
    // Implementation will use the transport to start the server
}

template<typename LifespanT, typename Func>
void FastMCPServer::lifespan(Func&& func) {
    static_assert(std::is_invocable_v<Func, LifespanT&>,
                 "Lifespan function must be invocable with LifespanT&");

    auto lifespan_mgr = std::make_unique<LifespanManager<LifespanT>>();
    lifespan_mgr->on_startup(std::forward<Func>(func));
    lifespan_manager_ = std::move(lifespan_mgr);
}

template<typename Func>
void FastMCPServer::tool(const std::string& name, Func&& func, const ToolOptions& options) {
    // Implementation will register the tool with the tool manager
}

template<typename Func>
void FastMCPServer::resource(const std::string& uri_template, Func&& func, const ResourceOptions& options) {
    // Implementation will register the resource with the resource manager
}

template<typename Func>
void FastMCPServer::prompt(const std::string& name, Func&& func, const PromptOptions& options) {
    // Implementation will register the prompt with the prompt manager
}

template<typename ServerSessionT, typename LifespanT>
Context<ServerSessionT, LifespanT>::Context(std::weak_ptr<ServerSessionT> session,
                                             std::shared_ptr<LifespanT> lifespan_context,
                                             types::RequestId request_id)
    : session_(std::move(session))
    , lifespan_context_(std::move(lifespan_context))
    , request_id_(std::move(request_id)) {}

template<typename ServerSessionT, typename LifespanT>
void Context<ServerSessionT, LifespanT>::debug(const std::string& message) {
    if (auto session = session_.lock()) {
        session->log(types::LoggingMessage{
            .level = types::LogLevel::Debug,
            .message = message
        });
    }
}

template<typename ServerSessionT, typename LifespanT>
void Context<ServerSessionT, LifespanT>::info(const std::string& message) {
    if (auto session = session_.lock()) {
        session->log(types::LoggingMessage{
            .level = types::LogLevel::Info,
            .message = message
        });
    }
}

template<typename ServerSessionT, typename LifespanT>
void Context<ServerSessionT, LifespanT>::warning(const std::string& message) {
    if (auto session = session_.lock()) {
        session->log(types::LoggingMessage{
            .level = types::LogLevel::Warning,
            .message = message
        });
    }
}

template<typename ServerSessionT, typename LifespanT>
void Context<ServerSessionT, LifespanT>::error(const std::string& message) {
    if (auto session = session_.lock()) {
        session->log(types::LoggingMessage{
            .level = types::LogLevel::Error,
            .message = message
        });
    }
}

template<typename ServerSessionT, typename LifespanT>
void Context<ServerSessionT, LifespanT>::report_progress(double progress,
                                                         std::optional<double> total,
                                                         std::optional<std::string> message) {
    if (auto session = session_.lock()) {
        session->send_progress(types::ProgressNotification{
            .token = request_id_,
            .progress = progress,
            .total = total,
            .message = std::move(message)
        });
    }
}

template<typename ServerSessionT, typename LifespanT>
Result<types::ResourceContent> Context<ServerSessionT, LifespanT>::read_resource(const std::string& uri) {
    if (auto session = session_.lock()) {
        return session->read_resource(uri);
    }
    return Error::internal_error("Session expired");
}

template<typename ServerSessionT, typename LifespanT>
template<typename SchemaT>
Result<types::ElicitationResult> Context<ServerSessionT, LifespanT>::elicit(const std::string& message, const SchemaT& schema) {
    if (auto session = session_.lock()) {
        return session->elicit(message, schema);
    }
    return Error::internal_error("Session expired");
}

template<typename ServerSessionT, typename LifespanT>
Result<types::SamplingResult> Context<ServerSessionT, LifespanT>::create_message(
    const std::vector<types::SamplingMessage>& messages,
    const types::SamplingOptions& options) {
    if (auto session = session_.lock()) {
        return session->create_message(messages, options);
    }
    return Error::internal_error("Session expired");
}

template<typename ServerSessionT, typename LifespanT>
ServerSessionT& Context<ServerSessionT, LifespanT>::session() {
    if (auto session = session_.lock()) {
        return *session;
    }
    throw std::runtime_error("Session expired");
}

template<typename ServerSessionT, typename LifespanT>
LifespanT& Context<ServerSessionT, LifespanT>::lifespan_context() {
    if (!lifespan_context_) {
        throw std::runtime_error("No lifespan context available");
    }
    return *lifespan_context_;
}

} // namespace server
} // namespace mcp